#本脚本实现的目的是为了将虚幻4中无法导出的cable组件信息，通过搜索的方法将其导出分割图
import argparse
from pathlib import Path
import cv2
import numpy as np

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument('--img-dir', default="", type=str)
    parser.add_argument('--seg-dir', default="", type=str)
    return parser.parse_args()


dx = [-1, 0, 1, 0]
dy = [0, 1, 0, -1]

def bfs(coordinates, origin_img, seg_img,color):
    h,w = seg_img.shape[:2]
    while len(coordinates) > 0:
        x,y = coordinates.pop(0)
        if origin_img[x][y][2] > 150 and origin_img[x][y][2] <= 255:
            seg_img[x][y] = np.array([255,0,0])
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if nx >= 0 and nx < h and ny >= 0 and ny < w and origin_img[x][y][2] > 200 and origin_img[x][y][2] <= 255 and (np.array_equal(seg_img[nx][ny], np.array([0,0,0]))) and ((nx,ny) not in coordinates):
                coordinates.append((nx, ny))

def get_fit_edge_maxpoint(seg_img, color):
    bin_img = cv2.inRange(seg_img,color, color)
    contours, hierarchy = cv2.findContours(bin_img,cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    rotateds = []
    new_contours = []
    coordinates = []
    for i in range(0,len(contours)):
        if(len(contours[i]) < 5):
            continue
        rotateds.append(cv2.fitEllipse(contours[i]))
        new_contours.append(contours[i])
    print(len(rotateds))
    idx = -1
    h_w_rate = -1
    for i in range(0, len(rotateds)):
        if(rotateds[i][1][0]<= 1e-2):
            continue
        rate = rotateds[i][1][1] / rotateds[i][1][0]
        if rate > h_w_rate:
            idx = i
            h_w_rate = rate
    for point in new_contours[idx]:
        x, y = point[0][1], point[0][0]
        coordinates.append((x, y))
    return coordinates

def get_fit_edge_point(img, color):
    bin_img = cv2.inRange(img,color, color)
    contours, hierarchy = cv2.findContours(bin_img,cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    coordinates = []
    for contour in contours:
        for point in contour:
            x, y = point[0][1], point[0][0]
            coordinates.append((x, y))
    return coordinates
def main():
    args = Parse_Arguments()
    origin_img = cv2.imread("imgs/ue4/origin_13.png")
    seg_img = cv2.imread("imgs/ue4/seg_13.png")
    seg_img_copy = seg_img.copy()
    dog_color = np.array([70, 52, 146]) #BGR格式
    human_color = np.array([143, 149, 226])
    rope_color = np.array([0,0,255])

    # 初始化坐标列表
    # coordinates = [(870,1641), (1250,1726)]
    coordinates = []
    coordinates = get_fit_edge_point(seg_img, dog_color)  #获取狗的边缘轮廓点集合
    coordinates += get_fit_edge_point(seg_img, human_color)  #获取人的边缘轮廓点集合
    bfs(coordinates,origin_img,seg_img,rope_color)  #第一次广搜
    coordinates += get_fit_edge_maxpoint(seg_img, np.array([255,0,0]))  #获取轮廓最大的框边缘
    bfs(coordinates, origin_img, seg_img_copy, rope_color)  #第二次广搜
    cv2.imwrite("111.png",seg_img)
    cv2.waitKey(0)
if __name__ == "__main__":
    main()
